from keras.layers import Dense
from keras.models import Model
from keras.callbacks import ModelCheckpoint
from keras.utils import to_categorical
from keras.metrics import top_k_categorical_accuracy
from keras.applications.inception_resnet_v2 import InceptionResNetV2,preprocess_input
import numpy as np
import os
import datetime
from data_process import myDataSet
from assumption import *
from classify import class_names_to_ids

"""指定训练or测试"""
TRAIN=0x00
TEST=0x01
NULL=0x02
# 指定任务类型
mission=TRAIN

"""训练"""
WEIGHTS=0X03
FULL_MODEL=0X04
# 指定模型保存类型
model_type=FULL_MODEL

# 是否使能last_save
last_save=True

# 保存模型的文件夹
save_dir="trained_models"

"""测试"""
# 指定测试模型名称
test_model_path=None


#################################################################


# 记录开始时间
start_time=datetime.datetime.now()

"""指定数据集信息"""
# 创建数据集对象,生成索引文件,并随机划分训练集、测试集、验证集索引
ds=myDataSet(class_names_to_ids,"datasets/trashnet/data/dataset-resized/dataset-resized/",
            "index/list.txt","index/list_train.txt","index/list_test.txt","index/list_val.txt",
            preprocess_input,new_index_file=True)


"""构建模型
1. 权重使用imagenet数据集预训练
2. 全局平均池化?
3. 不加载全连接层，按照需要分类的类别数(6)进行设置"""
base_model:Model = InceptionResNetV2(include_top=False, pooling='avg')
outputs = Dense(6, activation='softmax')(base_model.output)
model = Model(base_model.inputs, outputs)
print("模型构建完成")


"""训练模型
# 问题:
## 1. 验证用的是测试集合，是否合理?"""
if(mission==TRAIN):
    # 获取训练集 & 验证集
    x_train,y_train=ds.Load("train")
    x_test,y_test=ds.Load("val")    # 历史遗留问题，使用了val & test混用
    train_images=x_train
    test_images=x_test
    # 转换标签为one-hot编码
    y_train = to_categorical(y_train)
    y_test = to_categorical(y_test)
    print("数据集加载完成")

    # 设置ModelCheckpoint，按照验证集的准确率进行保存
    filepath="model_{epoch:02d}.hdf5" # 保存模型的名称格式: 模型名称_第几个epoch_验证集准确率
    checkpoint = ModelCheckpoint(os.path.join(save_dir, filepath), monitor='val_accuracy',verbose=1, 
                                save_best_only=True,save_weights_only=(model_type==WEIGHTS))
    if(model_type==FULL_MODEL):
        print("check point: 保存完整模型")
    elif(model_type==WEIGHTS):
        print("check point: 保存权重模型")

    # 使用adam优化器，交叉熵损失函数，设置监控指标为准确率,准确率指标为top3准确率、top5准确率
    model.compile(optimizer='adam',
                loss='categorical_crossentropy',
                metrics=['accuracy', acc_top3, acc_top5])
    print("模型编译完成,开始训练")

    # 模型训练(训练代数为5)
    model.fit(np.array(train_images), y_train,
            batch_size=16,
            epochs=50,
            shuffle=True,
            validation_data=(np.array(test_images), y_test),
            callbacks=[checkpoint])

    if(last_save==True):
        # 保存最后一代模型(以防丢失数据)
        for i in range(6):
            try:
                if(i==0):
                    # h5格式
                    model.save(os.path.join(save_dir, "last_model.hdf5"),save_format="h5")
                elif(i==1):
                    # tf格式,保存在last_model_tf文件夹下
                    model.save(os.path.join(save_dir, "last_model_tf"),save_format="tf")
                elif(i==2):
                    # keras格式，保存为单一keras文件
                    model.save(os.path.join(save_dir, "last_model.keras"),save_format="keras")
                elif(i==3):
                    # keras格式，保存在last_model_keras文件夹下
                    model.save(os.path.join(save_dir, "last_model_keras"),save_format="keras")
                elif(i==4):
                    # 保存权重模型(h5)
                    model.save_weights(os.path.join(save_dir, "last_weights.hdf5"),save_format="h5")
                elif(i==5):
                    # 保存权重模型(tf),保存在last_weights_tf文件夹下
                    model.save_weights(os.path.join(save_dir, "last_weights_tf"),save_format="tf")
                print("保存模型成功:{}".format(i))
            except Exception as e:
                print("保存模型失败:{}".format(i))
                print(e)
                pass
    

    """测试模型"""
elif(mission==TEST):
    # 获取测试集
    x_test,y_test=ds.Load("val")    # 历史遗留问题，使用了val作为测试集
    print("数据集加载完成")
    test_images=x_test

    # 加载指定模型
    model_path=os.path.join(save_dir,test_model_path)
    model.load_weights(model_path)  # 权重模型 .h5文件
    print("权重模型加载完成")

    # 直接使用predict方法进行预测
    y_pred = model.predict(np.array(test_images))
    print("predict result:\n", y_pred)
    
    # 计算准确率
    y_pred_class = np.argmax(y_pred, axis=1) # 在每一行中选取最大值的列索引作为预测类别
    y_true_class = y_test
    test_num = len(y_true_class)
    true_num=0
    for i in range(test_num):
        if y_pred_class[i] == y_true_class[i]:
            true_num+=1
    accuracy = float(true_num/test_num)
    print("accuracy: {}%".format(accuracy * 100))


duration=datetime.datetime.now()-start_time
print("训练或测试完成,总耗时:{}".format(duration))
